# 导入Python库
import matplotlib
# 本训练模型基于Caltech101 的图像分类， Caltech101 包括102类
# images路径存放图片
# train.txt存放训练集的图片
# class.txt存放分类图片的类型
# epoch_10为训练完成的模型，为了减少时间消耗，并未使用验证集，并且训练轮次设定为10轮
matplotlib.use('Agg')
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import paddlex as pdx
from paddlex import transforms as T

# 设置数据增强的方式
# API说明：https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/apis/transforms/transforms.md
train_transforms = T.Compose([
    T.RandomCrop(crop_size=224),
    T.RandomHorizontalFlip(),
    T.Normalize()
])

# 定义训练和验证所用的数据集
# 根据具体情况修改图片所在位置
# API说明：https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/apis/datasets.md
train_dataset = pdx.datasets.ImageNet(
    data_dir='data/dataset/images',
    file_list='data/dataset/train.txt',
    label_list='data/dataset/class.txt',
    transforms=train_transforms,
    shuffle=True)

# 初始化模型，并进行训练
num_classes = len(train_dataset.labels)
model = pdx.cls.ResNet50_vd_ssld(num_classes=num_classes)

# API说明：https://github.com/PaddlePaddle/PaddleX/blob/release/2.0.0/docs/apis/models/classification.md
# 各参数介绍与调整说明：https://github.com/PaddlePaddle/PaddleX/tree/release/2.0.0/docs/parameters.md
model.train(
    num_epochs=10,
    train_dataset=train_dataset,
    train_batch_size=64,
    lr_decay_epochs=[4, 6, 8],
    learning_rate=0.025,
    save_dir='output/resnet50',
    use_vdl=True)

import os

# test_files存放测试集的图片
test_files = []
with open('data/dataset/test.txt', 'r') as file_to_read:
    for line in file_to_read:
        test_files.append(os.path.join('data/dataset/images', line.strip()))

import paddlex as pdx

# 模型载入
model = pdx.load_model('output/resnet50/epoch_10')
result_list = model.predict(test_files, train_transforms)

with open('result.txt', mode='w') as file_out:
    k = 0
    for i in range(len(result_list)):
        img_name = test_files[k].split('/')[-1]
        file_out.write(img_name + '\t' + '%s' % (result_list[i][0].get('category_id')) + '\n')
        k += 1
